import torch
import os
from utils import make_project_dir,save_image
from loss_utils import LossWriter
# from UNet import UNet
from fcn import FCN
import torch.optim as optim
import torch.nn as nn
from datasets import SegDatasets


# 定义分割网络的训练和验证超参数
BETA1 = 0.9                  # 优化算法中的超参数
BETA2 = 0.999                # 优化算法中的超参数
DATA_ROOT = "human_dataset"  # 存储数据集的根目录
INPUT_DIR_NAME = "human"     # 数据集中存储输入图像的目录名称
LABEL_DIR_NAME = "mask"      # 数据集中存储标签（掩码）的目录名称
LR = 0.0001                  # 学习率，用于控制模型参数在训练过程中的更新步长
BATCH_SIZE = 8               # 训练时的批次大小
H_FLIP = True                # 布尔变量，是否进行水平翻转的数据增强
V_FLIP = True                # 布尔变量，是否进行垂直翻转的数据增强
RESULTS_DIR = "results"      # 存储训练结果的目录
EPOCHS = 200                  # 训练的总轮数 
IMAGE_SIZE = 224             # 输入图像的尺寸
IMG_SAVE_FREQ = 100          # 保存训练过程中分割结果的频率
PTH_SAVE_FREQ = 2            # 保存模型权重的频率

VAL_BATCH_SIZE = 1           # 验证时的批次大小
VAL_FREQ = 1                 # 验证的频率

# 设置使用 cuda 运行模型
device = torch.device("cuda")

# 构建训练和验证dataloader
train_dataset = SegDatasets(IMAGE_SIZE, DATA_ROOT, INPUT_DIR_NAME, LABEL_DIR_NAME, H_FLIP, V_FLIP, train=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=BATCH_SIZE,
                                           shuffle=True)
val_dataset = SegDatasets(IMAGE_SIZE, DATA_ROOT, INPUT_DIR_NAME, LABEL_DIR_NAME, H_FLIP, V_FLIP, train=False)
val_loader = torch.utils.data.DataLoader(dataset=val_dataset,
                                             batch_size=VAL_BATCH_SIZE,
                                             shuffle=True)


# 定义BCE损失
bce_func = nn.BCELoss()

# 定义分割网络，并将网络参数绑定到Adam优化器
seg_net = FCN(num_classes=1).to(device)
optimizer = optim.Adam(params=seg_net.parameters(),
                       lr=LR,
                       betas=(BETA1, BETA2))
# 创建项目目录结构
make_project_dir(RESULTS_DIR, RESULTS_DIR)
loss_writer = LossWriter(os.path.join(RESULTS_DIR, "loss"))

# 定义训练函数
def train():
    iteration = 0 # 迭代次数
    for epo in range(1, EPOCHS):
        # 遍历dataloader中所有的数据
        for data in train_loader:
            human = data["human"].to(device)
            mask = data["mask"].to(device)

            # 将人像图片human输入到分割网络中，获得预测的掩码mask
            predict_mask = seg_net(human)

            # 通过预测mask和真实mask计算损失值
            bce_loss = bce_func(predict_mask, mask)

            # 清空梯度，更新网络参数
            optimizer.zero_grad()
            bce_loss.backward()
            optimizer.step()

            # 记录分割损失值，并输出到控制台
            loss_writer.add("bce_loss", bce_loss.item(), iteration)

            print("Iter: {}, BCE Loss: {:.4f}".format(iteration,
                                              bce_loss.item()))
            # 更新迭代次数
            iteration += 1

            # 保存训练过程中的分割mask
            if iteration % IMG_SAVE_FREQ == 0: # 每隔IMG_SAVE_FREQ次迭代保存一次
                train_patch = torch.cat((predict_mask, mask), dim=3) # 拼接预测mask和真实mask
                save_image(train_patch[0],
                           out_name=os.path.join(RESULTS_DIR, "train_images",
                                                str(iteration) + ".png")) # 保存训练过程中的分割mask
        # 保存分割网络权重
        if epo % PTH_SAVE_FREQ == 0: # 每隔PTH_SAVE_FREQ次迭代保存一次
            torch.save(seg_net.state_dict(),
                       os.path.join(RESULTS_DIR, "pth", str(epo) + ".pth"),
                       _use_new_zipfile_serialization=False) # 保存分割网络权重

        # 计算验证集的分割效果
        if epo % VAL_FREQ == 0:
            seg_net.eval()
            with torch.no_grad():
                # 遍历验证集，并进行分割
                for data in val_loader:
                    human = data["human"].to(device) # 获取人像图像
                    mask = data["mask"].to(device)   # 获取人像掩码
                    img_name = data["img_name"]      # 获取图像名称
                    predict_mask = seg_net(human)    # 通过分割网络进行分割
                    val_patch = torch.cat((predict_mask, mask), dim=3) # 拼接预测mask和真实mask
                    # 保存验证集分割结果，用于观察模型表现
                    save_image(val_patch[0],
                               out_name=os.path.join(RESULTS_DIR, "val_images",
                                                      img_name[0]))
            seg_net.train()

if __name__ == "__main__":
    train()